import argparse

import torch
import torch.nn  as nn
import torch.nn.functional as F
import numpy as np
from utils import overlap_and_add
from mindspore import Tensor
import numpy as np
import mindspore
import torch
import mindspore.ops as ops

input_x = Tensor(np.zeros((4, 4, 4)), mindspore.int32)
indices = Tensor(np.array([[0], [2], [1]]), mindspore.int32)
updates = Tensor(np.array([[[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]],
                           [[1, 2, 3, 4], [5, 6, 7, 8], [1, 2, 3, 4], [8, 8, 8, 8]],
                           [[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]]]), mindspore.int32)
scatter_nd_add = ops.ScatterNdAdd()
print("'''''''''''''''''''''''''准备输出'''''''''''''''''''''''''''")
output = scatter_nd_add(input_x, indices, updates)

print(output)
print("'''''''''''''''''''''''''''''''")